
#include "RSAChecking.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe signature checking.
 * 
 * @param rsapubkey The RSA public key to use for signature
 * checking. */
RSAChecking::RSAChecking(const RSAPublicKey &rsapubkey)
{
   rsapubkey.getSSLKeyComponents(&key);
   input = NULL;
   output = NULL;
   final_data_length = 0;
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe signature checking.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param rsapubkey The RSA public key to use for signature
 * checking. */
RSAChecking::RSAChecking(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const RSAPublicKey &rsapubkey)
{
   rsapubkey.getSSLKeyComponents(&key);
   input = &iinput;
   output = &ioutput;
   final_data_length = 0;
}

/** Destructor for the RSA public key class for signature checking. */
RSAChecking::~RSAChecking()
{
   RSA_free(key);
}

#include "../utils/StringUtils.h"

/** Check the signature of the data block and return the unsigned
 * data. The new unsigned data that is returned must be freed by the
 * caller.
 * 
 * @param signedblock A block of data with a signature appended to it.
 * 
 * @param signedsize The length of the signed data block.
 * 
 * @param[out] datasize The length of the data block returned. The
 * data returned must be freed by the caller.
 * 
 * @return A newly copied data block without the signature
 * appended. */
char *RSAChecking::processDataBlock(const char *signedblock, apr_size_t signedsize, apr_size_t *datasize)
{
   if (signedblock == NULL)
   {
      errprint("tried to check NULL data block");
      return NULL;
   }
   unsigned int sign_length;
   unsigned int network_siglength = *((unsigned int *)(&(signedblock[signedsize-sizeof(unsigned int)])));
   ((unsigned char *)(&sign_length))[0] = ((network_siglength & 0xFF000000) >> 24);
   ((unsigned char *)(&sign_length))[1] = ((network_siglength & 0xFF0000) >> 16);
   ((unsigned char *)(&sign_length))[2] = ((network_siglength & 0xFF00) >> 8);
   ((unsigned char *)(&sign_length))[3] = (network_siglength & 0xFF);
   EVP_PKEY *newpkey = EVP_PKEY_new();
   if (newpkey == NULL)
   {
      errprint("could not initialise the public key");
      return NULL;
   }
   if (key)
      EVP_PKEY_set1_RSA(newpkey, key);
   unsigned int max_sign_length = (key) ? EVP_PKEY_size(newpkey) : (RSA_BITS / 8);
   // ensure it isn't too long if something goes wrong
   sign_length = (sign_length > max_sign_length) ? max_sign_length : sign_length;
   EVP_MD_CTX ctxt;
   EVP_MD_CTX_init(&ctxt);
   if (!EVP_VerifyInit_ex(&ctxt, EVP_sha256(), NULL))
   {
      errprint("could not initialise signing");
      return NULL;
   }
   if (!EVP_VerifyUpdate(&ctxt, (void *)signedblock, signedsize - (sign_length + sizeof(unsigned int))))
   {
      errprint("could not add data to signing");
      return NULL;
   }
   char *signature = (char *)malloc(sizeof(char)*sign_length);
   memcpy(signature, &(signedblock[signedsize - (sign_length + sizeof(unsigned int))]), sizeof(char)*sign_length);
   int verify_result = EVP_VerifyFinal(&ctxt, (unsigned char *)signature, sign_length, newpkey);
   verified = (verify_result == 1);
   if (datasize)
      *datasize = signedsize - (sign_length + sizeof(unsigned int));
   EVP_PKEY_free(newpkey);
   free(signature);
   char *final_data = (char *)malloc(sizeof(char)*(signedsize - (sign_length + sizeof(unsigned int))));
   memcpy(final_data, signedblock, sizeof(char)*(signedsize - (sign_length + sizeof(unsigned int))));
   EVP_MD_CTX_cleanup(&ctxt);
   return final_data;
}

/** Begin the process of checking the signature on a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAChecking::processPipeBegin()
{
   EVP_MD_CTX_init(&e_ctxt);
   if (!EVP_VerifyInit_ex(&e_ctxt, EVP_sha256(), NULL))
   {
      errprint("could not initialise message digest for signature verification");
      return 1;
   }
   unsigned int max_sign_length = (key) ? RSA_size(key) : (RSA_BITS / 8);
   final_data_length = (max_sign_length + sizeof(unsigned int));
   return 0;
}

/** Process any outstanding data on the input pipe, and place any
 * processed data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAChecking::processPipeStep()
{
   int input_block = input->realLength();
   // if there is less than what we want to leave on the pipe, leave
   // it on the pipe
   if (input_block <= final_data_length)
      return 0;
   input_block -= final_data_length;
   char *itemp_block = (char *)malloc(sizeof(char)*input_block);
   apr_size_t readbytes = input->readFromStream(itemp_block, input_block);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != input_block)
   {
      errprint("failed to read from stream (read %" APR_SIZE_T_FMT ", expected %" APR_SIZE_T_FMT ")", readbytes, (apr_size_t)input_block);
      free(itemp_block);
      return 1;
   }
   if (!EVP_VerifyUpdate(&e_ctxt, (void *)itemp_block, input_block))
   {
      errprint("could not add data to signing");
      free(itemp_block);
      return 1;
   }
   apr_size_t writebytes = output->writeToStream(itemp_block, input_block);
   if (writebytes != input_block)
   {
      errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(input_block));
      free(itemp_block);
      return 1;
   }
   input->removeFromStream(readbytes);
   free(itemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int RSAChecking::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding data in the pipe and finalise the
 * signature verification process.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAChecking::processPipeEnd()
{
   processPipeStep();
   int input_block = input->realLength();
   if (input_block != final_data_length)
   {
      errprint("remaining signature data is not expected length");
      return 1;
   }
   char *itemp_block = (char *)malloc(sizeof(char)*input_block);
   apr_size_t readbytes = input->readFromStream(itemp_block, input_block);
   unsigned int sign_length;
   unsigned int network_siglength = *((unsigned int *)(&(itemp_block[input_block - sizeof(unsigned int)])));
   ((unsigned char *)(&sign_length))[0] = ((network_siglength & 0xFF000000) >> 24);
   ((unsigned char *)(&sign_length))[1] = ((network_siglength & 0xFF0000) >> 16);
   ((unsigned char *)(&sign_length))[2] = ((network_siglength & 0xFF00) >> 8);
   ((unsigned char *)(&sign_length))[3] = (network_siglength & 0xFF);

   if (!EVP_VerifyUpdate(&e_ctxt, (void *)itemp_block, input_block - (sign_length + sizeof(unsigned int))))
   {
      errprint("could not add final data to signing");
      return 1;
   }
   int signverify = 0;
   int retval = 1;
   EVP_PKEY *newpkey = EVP_PKEY_new();
   if (newpkey == NULL)
   {
      errprint("could not initialise the private key");
      goto checkpipefinal;
   }
   if (key)
      EVP_PKEY_set1_RSA(newpkey, key);
   if (sign_length > input_block)
   {
      errprint("signature length is larger than the held final buffer");
      goto checkpipefinal;
   }
   signverify = EVP_VerifyFinal(&e_ctxt, (const unsigned char *)(&(itemp_block[input_block - (sign_length + sizeof(unsigned int))])), sign_length, newpkey);
   if (signverify == -1) SSLCheckError();
   output->writeToStream(itemp_block, input_block - (sign_length + sizeof(unsigned int)));
   retval = 0;
checkpipefinal:
   free(itemp_block);
   EVP_PKEY_free(newpkey);
   verified = (signverify == 1);
   EVP_MD_CTX_cleanup(&e_ctxt);
   return retval;
}
